import numpy as np
import abc
import sys
from .graph import default_graph


class Node():
    """
    计算图节点类基类
    """

    def __init__(self, *parents, **kwargs):
        # 计算图对象，默认为全局对象
        self.graph = kwargs.get('graph', default_graph)
        self.need_save = kwargs.get('need_save', True)
        self.gen_node_name(**kwargs)

        self.parents = list(parents)
        self.children = []  # 子节点列表
        self.value = None  # 本节点的值
        self.jacobi = None  # 结果节点对本节点的雅克比矩阵

        # 将本节点添加到父节点的子节点列表中
        for parent in self.parents:
            parent.children.append(self)

        # 将本节点添加到计算图中
        self.graph.add_node(self)

    @abc.abstractmethod
    def compute(self):
        """
        根据父节点的值计算本节点的值
        """

    @abc.abstractmethod
    def get_jacobi(self, parent):
        """
        计算本节点对某个父节点的雅克比矩阵
        """

    def get_parents(self):
        return self.parents

    def get_children(self):
        return self.children

    def gen_node_name(self, **kwargs):
        """
        生成的节点的名称
        """
        self.name = kwargs.get('name', f'{self.__class__.__name__}:{self.graph.node_count()}')
        if self.graph.name_scope:
            self.name = f'{self.graph.name_scope}/{self.name}'

    def forward(self):
        """
        前向传播计算本节点的值，若父节点的值未被计算，则递归调用父节点的forward函数
        """
        for node in self.parents:
            if node.value is None:
                node.forward()
        self.compute()

    def backward(self, result):
        """
        反向传播，计算结果节点对本节点的雅可比矩阵
        """
        if self.jacobi is None:
            if self is result:
                self.jacobi = np.mat(np.eye(self.dimension))
            else:
                self.jacobi = np.mat(np.zeros((result.dimension, self.dimension)))
                for child in self.get_children():
                    if child.value is not None:
                        self.jacobi += child.backward(result) * child.get_jacobi(self)
        return self.jacobi

    def clear_jacobi(self):
        """
        清空结果节点对本节点的雅克比矩阵
        """
        self.jacobi = None

    @property
    def dimension(self):
        """
        返回本节点的值展开成向量后的维度
        """
        return self.value.shape[0] * self.value.shape[1]

    @property
    def shape(self):
        return self.value.shape

    def reset_value(self, recursive=True):
        """
        重置本节点的值，并递归重置本节点的下游节点的值
        """
        self.value = None
        if recursive:
            for child in self.children:
                child.reset_value()


class Variable(Node):
    """
    变量节点（计算图中的叶子节点）
    """

    def __init__(self, dim, init=False, trainable=True, **kargs):
        super().__init__(**kargs)
        self.dim = dim

        if init:
            self.value = np.mat(np.random.normal(0, 0.001, self.dim))

        self.trainable = trainable

    def set_value(self, value):
        assert isinstance(value, np.matrix) and value.shape == self.dim
        self.reset_value()
        self.value = value
